import torch
import torch.nn.functional as F
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

from model import CLIPWrapper
from config import CLIP_CONFIG

def load_model(model_path=None):
    """加载模型"""
    if model_path is None:
        # 使用预训练模型
        model = CLIPWrapper(CLIP_CONFIG['model_name'])
        print("使用预训练CLIP模型")
    else:
        # 加载训练好的模型
        checkpoint = torch.load(model_path, map_location='cpu')
        model = CLIPWrapper(CLIP_CONFIG['model_name'])
        model.load_state_dict(checkpoint['model_state_dict'])
        print(f"加载训练好的模型: {model_path}")
    
    model.eval()
    return model

def image_text_similarity(model, image_path, texts):
    """计算图像和文本的相似度"""
    # 加载图像
    image = Image.open(image_path).convert('RGB')
    
    # 处理输入
    inputs = model.processor(
        text=texts,
        images=image,
        return_tensors="pt",
        padding=True
    )
    
    # 编码
    with torch.no_grad():
        image_features = model.encode_image(inputs['pixel_values'])
        text_features = model.encode_text(inputs['input_ids'])
    
    # 计算相似度
    similarity = model.compute_similarity(image_features, text_features)
    
    # 转换为概率
    probs = F.softmax(similarity, dim=-1)
    
    return probs.squeeze().numpy()

def text_text_similarity(model, query_text, candidate_texts):
    """计算文本和文本的相似度"""
    # 处理输入
    inputs = model.processor(
        text=[query_text] + candidate_texts,
        return_tensors="pt",
        padding=True
    )
    
    # 编码
    with torch.no_grad():
        text_features = model.encode_text(inputs['input_ids'])
    
    # 计算相似度
    query_features = text_features[0:1]
    candidate_features = text_features[1:]
    
    similarity = torch.matmul(query_features, candidate_features.T)
    probs = F.softmax(similarity, dim=-1)
    
    return probs.squeeze().numpy()

def text_image_similarity(model, query_text, image_paths):
    """计算文本和图像的相似度"""
    # 处理文本输入
    text_inputs = model.processor(
        text=[query_text],
        return_tensors="pt",
        padding=True
    )
    
    # 编码文本
    with torch.no_grad():
        text_features = model.encode_text(text_inputs['input_ids'])
    
    # 处理并编码所有图像
    image_features_list = []
    for image_path in image_paths:
        try:
            image = Image.open(image_path).convert('RGB')
            image_inputs = model.processor(
                images=image,
                return_tensors="pt"
            )
            
            with torch.no_grad():
                image_features = model.encode_image(image_inputs['pixel_values'])
            image_features_list.append(image_features)
        except Exception as e:
            print(f"处理图像 {image_path} 失败: {e}")
            # 添加零向量作为占位符
            image_features_list.append(torch.zeros_like(text_features))
    
    # 计算相似度
    similarities = []
    for image_features in image_features_list:
        similarity = torch.matmul(text_features, image_features.T)
        similarities.append(similarity.item())
    
    return similarities

def zero_shot_classification(model, image_path, class_names):
    """零样本分类"""
    # 为每个类别创建描述
    texts = [f"a photo of a {name}" for name in class_names]
    
    # 计算相似度
    probs = image_text_similarity(model, image_path, texts)
    
    # 获取预测结果
    predicted_class_idx = np.argmax(probs)
    predicted_class = class_names[predicted_class_idx]
    confidence = probs[predicted_class_idx]
    
    return predicted_class, confidence, probs

def visualize_similarity(image_path, texts, similarities):
    """可视化相似度结果"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # 显示图像
    image = Image.open(image_path)
    ax1.imshow(image)
    ax1.axis('off')
    ax1.set_title('Input Image')
    
    # 显示相似度条形图
    y_pos = np.arange(len(texts))
    ax2.barh(y_pos, similarities)
    ax2.set_yticks(y_pos)
    ax2.set_yticklabels(texts)
    ax2.set_xlabel('Similarity Score')
    ax2.set_title('Text-Image Similarity')
    
    plt.tight_layout()
    plt.show()

def visualize_retrieval(query_text, image_paths, similarities):
    """可视化检索结果"""
    # 按相似度排序
    sorted_indices = np.argsort(similarities)[::-1]
    top_k = min(5, len(image_paths))
    
    fig, axes = plt.subplots(1, top_k + 1, figsize=(15, 3))
    
    # 显示查询文本
    axes[0].text(0.5, 0.5, f'Query:\n{query_text}', 
                ha='center', va='center', fontsize=10, wrap=True)
    axes[0].axis('off')
    axes[0].set_title('Query Text')
    
    # 显示前k个最相似的图像
    for i, idx in enumerate(sorted_indices[:top_k]):
        try:
            image = Image.open(image_paths[idx])
            axes[i+1].imshow(image)
            axes[i+1].axis('off')
            axes[i+1].set_title(f'Rank {i+1}\nSim: {similarities[idx]:.3f}')
        except Exception as e:
            axes[i+1].text(0.5, 0.5, f'Error\n{image_paths[idx]}', 
                          ha='center', va='center', fontsize=8)
            axes[i+1].axis('off')
    
    plt.tight_layout()
    plt.show()

def demo_zero_shot_classification():
    """演示零样本分类"""
    print("=== 零样本分类演示 ===")
    
    # 加载模型
    model = load_model()
    
    # CIFAR-10类别名称
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck']
    
    print("使用CIFAR-10数据集进行零样本分类演示")
    print(f"类别: {class_names}")
    
    try:
        # 尝试加载CIFAR-10数据集
        from torchvision import datasets, transforms
        import tempfile
        import os
        
        # 数据预处理
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        
        # 加载测试集
        test_dataset = datasets.CIFAR10(
            root='./data/cifar10', 
            train=False, 
            download=False,  # 从缓存加载
            transform=transform
        )
        
        print(f"✓ 加载CIFAR-10测试集成功，样本数: {len(test_dataset)}")
        
        # 选择一些样本进行演示
        demo_samples = []
        for i in range(5):
            image, label = test_dataset[i]
            # 转换为PIL图像
            image_pil = transforms.ToPILImage()(image)
            
            # 保存到临时文件
            temp_dir = tempfile.mkdtemp(prefix='cifar_demo_')
            image_path = os.path.join(temp_dir, f"cifar10_sample_{i}.jpg")
            image_pil.save(image_path)
            
            demo_samples.append((image_path, class_names[label]))
        
        # 演示第一个样本
        image_path, true_label = demo_samples[0]
        
        predicted_class, confidence, all_probs = zero_shot_classification(
            model, image_path, class_names
        )
        
        print(f"\n演示样本: CIFAR-10测试集样本")
        print(f"真实类别: {true_label}")
        print(f"预测类别: {predicted_class}")
        print(f"置信度: {confidence:.4f}")
        print(f"预测正确: {predicted_class == true_label}")
        
        print("\n所有类别概率:")
        for name, prob in zip(class_names, all_probs):
            print(f"  {name}: {prob:.4f}")
        
        # 可视化结果
        texts = [f"a photo of a {name}" for name in class_names]
        visualize_similarity(image_path, texts, all_probs)
        
        # 测试更多样本
        print("\n=== 更多样本测试 ===")
        correct_predictions = 0
        total_predictions = len(demo_samples)
        
        for i, (image_path, true_label) in enumerate(demo_samples):
            predicted_class, confidence, all_probs = zero_shot_classification(
                model, image_path, class_names
            )
            
            is_correct = predicted_class == true_label
            if is_correct:
                correct_predictions += 1
            
            print(f"\n样本 {i+1}:")
            print(f"  真实类别: {true_label}")
            print(f"  预测类别: {predicted_class}")
            print(f"  置信度: {confidence:.4f}")
            print(f"  预测正确: {is_correct}")
        
        accuracy = correct_predictions / total_predictions
        print(f"\n测试准确率: {accuracy:.2f} ({correct_predictions}/{total_predictions})")
        
        print(f"\n临时文件保存在: {temp_dir}")
        
    except Exception as e:
        print(f"❌ CIFAR-10演示失败: {e}")
        print("切换到几何图像演示模式...")
        
        # 备用演示：使用几何图像
        class_names = ['cat', 'dog', 'car', 'tree', 'building', 'person', 'bird', 'horse', 'ship', 'truck']
        image_path = 'data/images/animal_cat.jpg'
        
        try:
            predicted_class, confidence, all_probs = zero_shot_classification(
                model, image_path, class_names
            )
            
            print(f"预测类别: {predicted_class}")
            print(f"置信度: {confidence:.4f}")
            print("\n所有类别概率:")
            for name, prob in zip(class_names, all_probs):
                print(f"  {name}: {prob:.4f}")
            
            texts = [f"a photo of a {name}" for name in class_names]
            visualize_similarity(image_path, texts, all_probs)
            
        except Exception as e2:
            print(f"几何图像演示也失败: {e2}")

def demo_text_similarity():
    """演示文本相似度计算"""
    print("=== 文本相似度演示 ===")
    
    # 加载模型
    model = load_model()
    
    # CIFAR-10类别对应的文本描述
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck']
    texts = [f"a photo of a {name}" for name in class_names]
    
    print("使用CIFAR-10数据集进行文本相似度演示")
    print(f"文本描述: {texts}")
    
    try:
        # 尝试加载CIFAR-10数据集
        from torchvision import datasets, transforms
        import tempfile
        import os
        
        # 数据预处理
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        
        # 加载测试集
        test_dataset = datasets.CIFAR10(
            root='./data/cifar10', 
            train=False, 
            download=False,  # 从缓存加载
            transform=transform
        )
        
        print(f"✓ 加载CIFAR-10测试集成功，样本数: {len(test_dataset)}")
        
        # 选择一些样本进行演示
        demo_samples = []
        for i in range(3):
            image, label = test_dataset[i]
            # 转换为PIL图像
            image_pil = transforms.ToPILImage()(image)
            
            # 保存到临时文件
            temp_dir = tempfile.mkdtemp(prefix='cifar_similarity_')
            image_path = os.path.join(temp_dir, f"cifar10_sample_{i}.jpg")
            image_pil.save(image_path)
            
            demo_samples.append((image_path, class_names[label], texts[label]))
        
        # 演示第一个样本
        image_path, true_label, expected_text = demo_samples[0]
        
        similarities = image_text_similarity(model, image_path, texts)
        
        print(f"\n演示样本: CIFAR-10测试集样本")
        print(f"真实类别: {true_label}")
        print(f"期望文本: {expected_text}")
        
        # 找到最相似的文本
        max_similarity = np.max(similarities)
        max_index = np.argmax(similarities)
        predicted_text = texts[max_index]
        predicted_label = class_names[max_index]
        
        print(f"最相似文本: {predicted_text}")
        print(f"相似度: {max_similarity:.4f}")
        print(f"匹配正确: {predicted_label == true_label}")
        
        print("\n所有文本相似度:")
        for text, similarity in zip(texts, similarities):
            print(f"  {text}: {similarity:.4f}")
        
        # 可视化
        visualize_similarity(image_path, texts, similarities)
        
        # 测试更多样本
        print("\n=== 更多样本测试 ===")
        correct_matches = 0
        total_matches = len(demo_samples)
        
        for i, (image_path, true_label, expected_text) in enumerate(demo_samples):
            similarities = image_text_similarity(model, image_path, texts)
            
            # 找到最相似的文本
            max_similarity = np.max(similarities)
            max_index = np.argmax(similarities)
            predicted_text = texts[max_index]
            predicted_label = class_names[max_index]
            
            is_correct = predicted_label == true_label
            if is_correct:
                correct_matches += 1
            
            print(f"\n样本 {i+1}:")
            print(f"  真实类别: {true_label}")
            print(f"  期望文本: {expected_text}")
            print(f"  最相似文本: {predicted_text}")
            print(f"  相似度: {max_similarity:.4f}")
            print(f"  匹配正确: {is_correct}")
        
        accuracy = correct_matches / total_matches
        print(f"\n匹配准确率: {accuracy:.2f} ({correct_matches}/{total_matches})")
        
        print(f"\n临时文件保存在: {temp_dir}")
        
    except Exception as e:
        print(f"❌ CIFAR-10演示失败: {e}")
        print("切换到几何图像演示模式...")
        
        # 备用演示：使用几何图像
        texts = [
            "a photo of a cat",
            "a photo of a dog", 
            "a photo of a car",
            "a photo of a tree",
            "a photo of a building"
        ]
        image_path = 'data/images/animal_cat.jpg'
        
        try:
            similarities = image_text_similarity(model, image_path, texts)
            
            print(f"图像: {image_path}")
            print("\n文本相似度:")
            for text, similarity in zip(texts, similarities):
                print(f"  {text}: {similarity:.4f}")
            
            # 可视化
            visualize_similarity(image_path, texts, similarities)
            
        except Exception as e2:
            print(f"几何图像演示也失败: {e2}")

def demo_image_retrieval():
    """演示图像检索"""
    print("=== 图像检索演示 ===")
    
    # 加载模型
    model = load_model()
    
    # CIFAR-10类别对应的文本描述
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck']
    
    print("使用CIFAR-10数据集进行图像检索演示")
    print(f"可用类别: {class_names}")
    
    try:
        # 尝试加载CIFAR-10数据集
        from torchvision import datasets, transforms
        import tempfile
        import os
        
        # 数据预处理
        transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])
        
        # 加载测试集
        test_dataset = datasets.CIFAR10(
            root='./data/cifar10', 
            train=False, 
            download=False,  # 从缓存加载
            transform=transform
        )
        
        print(f"✓ 加载CIFAR-10测试集成功，样本数: {len(test_dataset)}")
        
        # 创建图像库（从测试集中选择一些样本）
        image_library = []
        temp_dir = tempfile.mkdtemp(prefix='cifar_retrieval_')
        
        # 每个类别选择2个样本
        samples_per_class = 2
        class_counts = {name: 0 for name in class_names}
        
        for i in range(len(test_dataset)):
            if len(image_library) >= len(class_names) * samples_per_class:
                break
                
            image, label = test_dataset[i]
            class_name = class_names[label]
            
            if class_counts[class_name] < samples_per_class:
                # 转换为PIL图像
                image_pil = transforms.ToPILImage()(image)
                
                # 保存到临时文件
                image_path = os.path.join(temp_dir, f"{class_name}_{class_counts[class_name]}.jpg")
                image_pil.save(image_path)
                
                image_library.append((image_path, class_name))
                class_counts[class_name] += 1
        
        print(f"✓ 创建图像库成功，包含 {len(image_library)} 张图像")
        
        # 演示图像检索
        query_text = "a photo of a cat"
        image_paths = [item[0] for item in image_library]
        
        similarities = text_image_similarity(model, query_text, image_paths)
        
        print(f"\n查询文本: '{query_text}'")
        print("\n图像相似度:")
        for (path, label), similarity in zip(image_library, similarities):
            print(f"  {label}: {similarity:.4f}")
        
        # 找到最相似的图像
        max_similarity = np.max(similarities)
        max_index = np.argmax(similarities)
        best_image_path, best_label = image_library[max_index]
        
        print(f"\n最相似图像: {best_label}")
        print(f"相似度: {max_similarity:.4f}")
        print(f"检索正确: {'cat' in best_label}")
        
        # 可视化
        visualize_retrieval(query_text, image_paths, similarities)
        
        # 测试多个查询文本
        print("\n=== 更多查询测试 ===")
        test_queries = [
            "a photo of an airplane",
            "a photo of a dog",
            "a photo of a ship"
        ]
        
        for query_text in test_queries:
            similarities = text_image_similarity(model, query_text, image_paths)
            
            # 找到最相似的图像
            max_similarity = max(similarities)
            max_index = similarities.index(max_similarity)
            best_image_path, best_label = image_library[max_index]
            
            expected_class = query_text.split()[-1]  # 获取查询中的类别
            is_correct = expected_class in best_label
            
            print(f"\n查询文本: '{query_text}'")
            print(f"最相似图像: {best_label}")
            print(f"相似度: {max_similarity:.4f}")
            print(f"检索正确: {is_correct}")
        
        print(f"\n临时文件保存在: {temp_dir}")
        
    except Exception as e:
        print(f"❌ CIFAR-10演示失败: {e}")

def main():
    """主演示函数"""
    print("CLIP模型演示")
    print("=" * 50)
    
    # 演示零样本分类
    demo_zero_shot_classification()
    
    # 演示文本相似度
    demo_text_similarity()
    
    # 演示图像检索
    demo_image_retrieval()
    
    print("\n演示完成!")

if __name__ == "__main__":
    main()